package com.tuniu.agents.prompt;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.ResourcePatternResolver;

import java.io.IOException;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * Implementation of the prompt manager
 */
public class ResourcePromptManager implements InitializingBean, PromptManager {

    private static final Logger logger = LoggerFactory.getLogger(ResourcePromptManager.class);

    private static final String TEMPLATE_INCLUDE_PATTERN = "\\{\\{include '(.+?)'}}";

    private final ResourcePatternResolver resourcePatternResolver;
    private final Map<String, String> templates = new HashMap<>();

    public ResourcePromptManager(ResourcePatternResolver resourcePatternResolver) {
        this.resourcePatternResolver = resourcePatternResolver;
    }

    @Value("${aigc.agents.promptLoaderPaths:classpath:/prompt/**/*.st}")
    private String promptLoaderPaths;

    @Override
    public void afterPropertiesSet() throws IOException {
        // Load all .st files in the resources/prompts directory
        String[] patterns = promptLoaderPaths.split(",");
        for (String pattern : patterns) {
            Resource[] resources = resourcePatternResolver.getResources(pattern);
            for (Resource resource : resources) {
                String filename = resource.getFilename();
                if (filename != null) {
                    String templateName = filename.replace(".st", "");
                    String templateContent = readResourceContent(resource);
                    templates.put(templateName, templateContent);
                }
            }
        }

        // Template content replacement: combine the content of different templates
        for (Map.Entry<String, String> entry : templates.entrySet()) {
            entry.setValue(renderTemplate(entry.getKey(), templates, new HashMap<>()));
        }

        logger.info("Resource Prompts List {}", templates.keySet());
    }

    private static String renderTemplate(String templateName, Map<String, String> templates, Map<String, String> renderedTemplates) {
        if (renderedTemplates.containsKey(templateName)) {
            return renderedTemplates.get(templateName);
        }

        String templateContent = templates.get(templateName);
        if (templateContent == null) {
            throw new IllegalArgumentException("Template not found: " + templateName);
        }

        Pattern pattern = Pattern.compile(TEMPLATE_INCLUDE_PATTERN);
        Matcher matcher = pattern.matcher(templateContent);

        StringBuffer renderedContent = new StringBuffer();

        while (matcher.find()) {
            String includedTemplateName = matcher.group(1);
            String includedContent = renderTemplate(includedTemplateName, templates, renderedTemplates);
            matcher.appendReplacement(renderedContent, includedContent);
        }
        matcher.appendTail(renderedContent);

        String finalRenderedContent = renderedContent.toString();
        renderedTemplates.put(templateName, finalRenderedContent);
        return finalRenderedContent;
    }

    private String readResourceContent(Resource resource) throws IOException {
        return resource.getContentAsString(Charset.defaultCharset());
    }

    public String getTemplateContent(String templateName) {
        String content = templates.get(templateName);
        if (content == null) {
            throw new IllegalArgumentException("Template not found: " + templateName);
        }
        return content;
    }

    public static String renderTemplate(String templateContent, Map<String, String> params) {
        for (String key : params.keySet()) {
            String placeholder = "{{" + key + "}}";
            if (templateContent.contains(placeholder)) {
                templateContent = templateContent.replace(placeholder, params.get(key));
            }
        }
        return templateContent;
    }

}
